#ifndef MULTIPLE_SIGN_TRACKER_H_
#define MULTIPLE_SIGN_TRACKER_H_

#include "opencv2/imgproc/imgproc.hpp"
#include "base/parallel_utility.hpp"
#include "base/km_matcher.hpp"
#include "local_sign.hpp"
#include "sign_tracker.hpp"
#include <map>

class MST 
{
private:
    std::list<Sign_Tracker*> ptrackers_;
    std::vector<LOCAL_SIGN>* pdetected_signs_;
    cv::Mat src_;

    const float kIOU_threshold_ = 0.8;
    const float kOLR_threshold_ = 0.8;
    float geometric_CRITERIA;

    
    /* debug related */
    size_t frame_index_ = 0;
    bool debug_message_ = true;

public:
    std::vector<std::vector<LOCAL_SIGN>> shaped_lists_;
    std::map<std::size_t, std::vector<LOCAL_SIGN>> pSt_lists_;

    MST(bool debug_message = false) : debug_message_(debug_message) {}
    virtual ~MST() { delete_all_tracker(); }

    void run(std::vector<LOCAL_SIGN> &detected_signs, const cv::Mat &src) 
    {
        src_ = src.clone();
        geometric_CRITERIA = src_.cols / 6;
        pdetected_signs_ = &detected_signs;
        assign_detections();
        delete_inactive_tracker();

        if (debug_message_) {
            std::cout << "Frame " << std::to_string(frame_index_) << " exists " << std::to_string(ptrackers_.size()) << " trackers..." << std::endl;
            std::cout << "---------------------------------------------------------" << std::endl << std::endl;
        }
        frame_index_++;
    }

    void cal_KM_cost(std::vector<std::vector<float>> &_KM_cost)
    {
        size_t sign_num = _KM_cost.size();
        // calculate the cos of each pair
        for (size_t i = 0; i < sign_num; ++i) {
            auto &sign = (*pdetected_signs_)[i];
            size_t j = 0;
            for (auto &&ptracker: ptrackers_) {
                if (debug_message_) {
                    std::cout << "[S_" << sign.id_ << ", T_" << ptracker->get_id() << "]: " << std::endl;
                }
                float apparance_cost = 0;
                float geometrical_cost = LfUtils::center_diff(sign.calibrated_bbox_, ptracker->get_bbox());
                if (geometrical_cost > geometric_CRITERIA) {
                    _KM_cost[i][j] = geometric_CRITERIA;
                } else {
                    apparance_cost = ptracker->detectAt(src_, sign.calibrated_bbox_);
                    _KM_cost[i][j] = geometrical_cost / apparance_cost;
                }
                    
                if (debug_message_)
                {
                    std::cout << "  geometric cost: " << geometrical_cost << std::endl;
                    std::cout << "  apparance cost:   " << apparance_cost << std::endl;
                    std::cout << "  KM_cost :   " << _KM_cost[i][j] << std::endl << std::endl;
                }
                j++;
            }
        }
    }

    void cal_KM_cost_parallel(std::vector<std::vector<float>> &_KM_cost)
    {
        size_t sign_num = _KM_cost.size();
        // calculate the cos of each pair
        for (size_t sign_idx = 0; sign_idx < sign_num; ++sign_idx) {
            auto &sign = (*pdetected_signs_)[sign_idx];
 
            auto do_work = [this, &sign, &sign_idx, &_KM_cost](size_t tracker_idx, std::mutex& _m) {
                auto itor_ptrackers = ptrackers_.begin();
                std::advance(itor_ptrackers, tracker_idx);
                auto &ptracker = (*itor_ptrackers);

                // if (debug_message_) {
                //     std::cout << "[S_" << sign.id_ << ", T_" << ptracker->get_id() << "]: " << std::endl;
                // }
                float apparance_cost = 0;
                float geometrical_cost = LfUtils::center_diff(sign.calibrated_bbox_, ptracker->get_bbox());
                if (geometrical_cost > geometric_CRITERIA) {
                    _KM_cost[sign_idx][tracker_idx] = geometric_CRITERIA;
                } else {
                    apparance_cost = ptracker->detectAt(src_, sign.calibrated_bbox_);
                    _KM_cost[sign_idx][tracker_idx] = geometrical_cost / apparance_cost;
                }
                    
                // if (debug_message_) {
                //     std::cout << "  geometric cost: " << geometrical_cost << std::endl;
                //     std::cout << "  apparance cost:   " << apparance_cost << std::endl;
                //     std::cout << "  KM_cost :   " << _KM_cost[sign_idx][tracker_idx] << std::endl << std::endl;
                // }
            };

        
            /* do_work parallelly */
            size_t chunk = LOOP_PARALISM::get_chunk(ptrackers_.size());
            for (size_t i = 0; i < chunk; ++i) {
                {
                    LOOP_PARALISM loop_paraller;
                    auto &j = loop_paraller.thread_index_;
                    for (j = 0; j < loop_paraller.numthreads_; ++j) {
                        auto &&index = i * chunk + j;
                        if (index >= ptrackers_.size()) { continue; }
                        loop_paraller.works_[j] = std::thread(do_work, index, std::ref(loop_paraller.m_));    
                    }
                }
            }
        }
    }

 
    void tracker_update(std::vector<int> &tracker_match_pair, std::vector<int> &sign_match_pair) {
        /* Behavior determination of existing tracker */
        for (size_t j = 0; j < tracker_match_pair.size(); ++j) {
            int sign_idx = tracker_match_pair[j];
            auto itor_ptrackers = ptrackers_.begin();
            std::advance(itor_ptrackers, j);
            auto &ptracker = (*itor_ptrackers);

            if (sign_idx != -1) {
                auto psign = &((*pdetected_signs_)[sign_idx]);
                if (debug_message_) { std::cout <<  "sign_" << psign->id_ << " "; }
                if(ptracker->update_tracker(src_, psign->calibrated_bbox_)) {
                    psign->corrsponding_bbox_ = ptracker->get_bbox();
                    psign->global_id_ = ptracker->get_id();
                    pSt_lists_[psign->global_id_].push_back(*psign);
                    psign->tracker_ = ptracker;
                    psign->debug_color_ = GREEN;
                } else {
                    ptracker->lose_tracker();
                    sign_match_pair[sign_idx] = -1;
                }
            } else {
                ptracker->lose_tracker();
            }
        }
    }

    void tracker_update_parallel(std::vector<int> &tracker_match_pair, std::vector<int> &sign_match_pair) {

        auto do_work = [this, &tracker_match_pair, &sign_match_pair](size_t tracker_idx, std::mutex& _m){
            int sign_idx = tracker_match_pair[tracker_idx];
            auto itor_ptrackers = ptrackers_.begin();
            std::advance(itor_ptrackers, tracker_idx);
            auto &ptracker = (*itor_ptrackers);

            if (sign_idx != -1) {
                auto psign = &((*pdetected_signs_)[sign_idx]);
                // if (debug_message_) { std::cout <<  "sign_" << psign->id_ << " "; }
                if(ptracker->update_tracker(src_, psign->calibrated_bbox_)) {
                    psign->corrsponding_bbox_ = ptracker->get_bbox();
                    psign->global_id_ = ptracker->get_id();
                    pSt_lists_[psign->global_id_].push_back(*psign);
                    psign->debug_color_ = GREEN;
                } else {
                    ptracker->lose_tracker();
                    sign_match_pair[sign_idx] = -1;
                }
            } else {
                ptracker->lose_tracker();
            }
        };

        /* do_work parallelly */
        size_t chunk = LOOP_PARALISM::get_chunk(tracker_match_pair.size());
        for (size_t i = 0; i < chunk; ++i) {
            {
                LOOP_PARALISM loop_paraller;
                auto &j = loop_paraller.thread_index_;
                for (j = 0; j < loop_paraller.numthreads_; ++j) {
                    auto &&index = i * chunk + j;
                    if (index >= ptrackers_.size()) { continue; }
                    loop_paraller.works_[j] = std::thread(do_work, index, std::ref(loop_paraller.m_));    
                }
            }
        }
    }


    /* acquire optimal matches */
    void assign_detections() 
    {
        const size_t sign_num = (*pdetected_signs_).size();
        const size_t tracker_num = ptrackers_.size();

        if (sign_num == 0 && tracker_num == 0) { return; }

        std::vector<int> sign_match_pair(sign_num, -1);
        std::vector<int> tracker_match_pair(tracker_num, -1);

        if (tracker_num > 0 && sign_num > 0)
        {
            /* calcuate the matching cost */
            std::vector<std::vector<float>> KM_cost(sign_num, std::vector<float>(tracker_num,0));
            auto get_KM_data = [&KM_cost](unsigned sign_idx, unsigned tracker_idx) 
            {
		        return KM_cost[sign_idx][tracker_idx];
	        };

            // calculate the cos of each pair
            // cal_KM_cost(KM_cost);
            cal_KM_cost_parallel(KM_cost);

            auto get_KM = std::bind(get_KM_data, std::placeholders::_1, std::placeholders::_2);
	        algo_utils::MunkresData<float>	KM_matching(sign_num, tracker_num, get_KM);
	        auto matches = KM_matching.solve();
            
            if (debug_message_) { std::cout << "------ Matches Info ------" << std::endl; }
	        for (const auto &match: matches) {
                auto &sign = (*pdetected_signs_)[match.first];
                auto itor_ptrackers = ptrackers_.begin();
                std::advance(itor_ptrackers, match.second);
                auto &ptracker = (*itor_ptrackers);

                // bool geometric_check = LfUtils::center_diff(sign.calibrated_bbox_, ptracker->get_bbox()) < geometric_CRITERIA;
                if ( KM_cost[match.first][match.second] != geometric_CRITERIA) {
                    sign_match_pair[match.first] = match.second;
                    tracker_match_pair[match.second] = match.first;
                    if (debug_message_) { std::cout << "  sign_" << sign.id_ << " matches tracker_" << ptracker->get_id() << std::endl; }
                } else {
                    if (debug_message_) { std::cout << "  sign_" << sign.id_ << " failed geometric check of tracker_" << ptracker->get_id() << std::endl; }
                }
	        }
            if (debug_message_) { std::cout << "--------------------------" << std::endl << std::endl; }
        }

        /* Behavior determination of existing tracker */
        tracker_update(tracker_match_pair, sign_match_pair);
        // tracker_update_parallel(tracker_match_pair, sign_match_pair);
   
        /* add new tracker for unassigned sign */
        for (size_t i = 0; i < sign_num; ++i) 
        {
            int tracker_idx = sign_match_pair[i];
            if (tracker_idx == -1)
            {
                auto psign = &((*pdetected_signs_)[i]);

                Sign_Tracker *new_tracker = new Sign_Tracker(debug_message_);
                if (debug_message_) { std::cout << "sign_" << psign->id_ << " "; }
                new_tracker->init(src_, psign->calibrated_bbox_);
                ptrackers_.push_back(new_tracker);
                psign->corrsponding_bbox_ = psign->calibrated_bbox_;
                psign->global_id_ = new_tracker->get_id();
                psign->tracker_ = new_tracker;
                pSt_lists_[psign->global_id_].push_back(*psign);
                psign->debug_color_ = BLUE;
            }
        }
    }

    void draw_result(cv::Mat &debug_img) 
    {
        int count = 0;
        for (auto &&ptracker: ptrackers_) 
        {
            auto color = MAGENTA;
            ptracker->draw_result(debug_img, color);
            std::string info = ptracker->info();
            cv::Point Text_pt(30, 40*count + 50);
            cv::putText(debug_img, info, Text_pt, cv::FONT_HERSHEY_COMPLEX, 0.8, color,2);
            count++;
        }
    }

private:
    void delete_inactive_tracker()
    {
        auto pp_tracker = ptrackers_.begin();
        while (pp_tracker != ptrackers_.end())
        {
            if ((*pp_tracker)->is_inactive())
            {
                std::size_t id = (*pp_tracker)->get_id();
                shaped_lists_.push_back(pSt_lists_[id]);
                std::map<std::size_t, std::vector<LOCAL_SIGN>>::iterator key = pSt_lists_.find(id);
                if(key != pSt_lists_.end())
                {
	                pSt_lists_.erase(key);
                }
                delete *pp_tracker;
                *pp_tracker = nullptr;
                ptrackers_.erase(pp_tracker++);
            } 
            else
            {
                pp_tracker++;
            }
        }
    }

    void delete_all_tracker() 
    {
        auto pp_tracker = ptrackers_.begin();
        while (pp_tracker != ptrackers_.end()) {
            delete *pp_tracker;
            *pp_tracker = nullptr;
            pp_tracker++;
        }
        ptrackers_.clear();
    }
};


#endif // MULTIPLE_SIGN_TRACKER_H_